library(ggplot2)
library(patchwork)
library(foreign)

# set to desired working directory
setwd("4_IAT")

# Run Models
source("10_IAT_fullsample.R")
source("10_IAT_blackwhite.R")
source("10_IAT_ideowhite.R")

# Start log file for script -----------------------------------------------
TeachingDemos::txtStart("./log files/figure4_data_export.txt")

# Initial Plot to Check ---------------------------------------------------
colors <- c("navy", "darkorange1")

p_full <- ggplot(p_dat_Fsamp, aes(x = week, y = pred, group = year)) +
  geom_vline(color = "red", aes(xintercept = as.Date("2020-5-22"))) +
  geom_linerange(aes(ymin = pred - 1.39*se,
                     ymax = pred + 1.39*se,
                     color = year),
                 position = position_dodge(width = .5),
                 lwd = .5) +
  geom_point(size = 2, aes(shape = year, color = year),
             position = position_dodge(width = .5)) +
  geom_smooth(data = subset(p_dat_Fsamp, week < "2020-5-25"), 
              span = .8,
              aes(x = week, y = pred, color = year, linetype = year), se = FALSE) +
  geom_smooth(data = subset(p_dat_Fsamp, week >= "2020-5-25"), 
              span = .8,
              aes(x = week, y = pred, color = year, linetype = year), se = FALSE) +
  theme_bw() +
  scale_linetype_manual(values = c("dotdash", "solid"), guide = FALSE) +
  scale_color_manual(values = colors,name = "") +
  scale_shape_manual(values = c(1, 16), name = "") +
  labs(y = "Predicted D-Scores",
       x = "", title = "Full Sample") +
  scale_x_date(breaks = unique(p_dat_Fsamp$week)[c(1, 6, 10, 15, 19, 23, 28, 32, 37, 41, 45, 50)],
               labels = c("Jan", "Feb", "Mar", "Apr", "May", "Jun",
                          "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"),
               minor_breaks = NULL) + 
  scale_y_continuous(labels = numform::ff_num(zero = 0, digits = 2),
                     limits = c(.2, .5)) +
  theme(legend.position = "bottom",
        plot.title = element_text(size = 20, hjust = 0.5),
        axis.title = element_text(size = 14),
        axis.text = element_text(size = 12),
        panel.grid.major = element_blank(), 
        panel.grid.minor = element_blank(),
        legend.text = element_text(size = 14),
        strip.text = element_text(size = 16),
        strip.background = element_blank())

p_wht <- ggplot(subset(p_dat_bw, race == "White"), aes(x = week, y = pred, group = year)) +
  geom_vline(color = "red", aes(xintercept = as.Date("2020-5-22"))) +
  geom_linerange(aes(ymin = pred - 1.39*se,
                     ymax = pred + 1.39*se,
                     color = year),
                 position = position_dodge(width = .5),
                 lwd = .5) +
  geom_point(size = 2, aes(shape = year, color = year),
             position = position_dodge(width = .5)) +
  geom_smooth(data = subset(p_dat_bw, race == "White" & week < "2020-5-25"), 
              span = .8,
              aes(x = week, y = pred, color = year, linetype = year), se = FALSE) +
  geom_smooth(data = subset(p_dat_bw, race == "White" & week >= "2020-5-25"), 
              span = .8,
              aes(x = week, y = pred, color = year, linetype = year), se = FALSE) +
  theme_bw() +
  scale_linetype_manual(values = c("dotdash", "solid"), guide = FALSE) +
  scale_color_manual(values = colors,name = "") +
  scale_shape_manual(values = c(1, 16), name = "") +
  labs(y = "Predicted D-Scores",
       x = "", title = "Among Whites") +
  scale_x_date(breaks = unique(subset(p_dat_bw, race == "White")$week)[c(1, 6, 10, 15, 19, 23, 28, 32, 37, 41, 45, 50)],
               labels = c("Jan", "Feb", "Mar", "Apr", "May", "Jun",
                          "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"),
               minor_breaks = NULL) + 
  scale_y_continuous(labels = numform::ff_num(zero = 0, digits = 2),
                     limits = c(.2, .5)) +
  theme(legend.position = "bottom",
        plot.title = element_text(size = 20, hjust = 0.5),
        axis.title = element_text(size = 14),
        axis.text = element_text(size = 12),
        panel.grid.major = element_blank(), 
        panel.grid.minor = element_blank(),
        legend.text = element_text(size = 14),
        strip.text = element_text(size = 20),
        strip.background = element_blank())


p_blk <- ggplot(subset(p_dat_bw, race == "Black"), aes(x = week, y = pred, group = year)) +
  geom_vline(color = "red", aes(xintercept = as.Date("2020-5-22"))) +
  geom_linerange(aes(ymin = pred - 1.39*se,
                     ymax = pred + 1.39*se,
                     color = year),
                 position = position_dodge(width = .5),
                 lwd = .5) +
  geom_point(size = 2, aes(shape = year, color = year),
             position = position_dodge(width = .5)) +
  geom_smooth(data = subset(p_dat_bw, race == "Black" & week < "2020-5-25"), 
              span = .8,
              aes(x = week, y = pred, color = year, linetype = year), se = FALSE) +
  geom_smooth(data = subset(p_dat_bw, race == "Black" & week >= "2020-5-25"), 
              span = .8,
              aes(x = week, y = pred, color = year, linetype = year), se = FALSE) +
  theme_bw() +
  scale_linetype_manual(values = c("dotdash", "solid"), guide = FALSE) +
  scale_color_manual(values = colors,name = "") +
  scale_shape_manual(values = c(1, 16), name = "") +
  labs(y = "Predicted D-Scores",
       x = "", title = "Among Blacks") +
  scale_x_date(breaks = unique(subset(p_dat_bw, race == "Black")$week)[c(1, 6, 10, 15, 19, 23, 28, 32, 37, 41, 45, 50)],
               labels = c("Jan", "Feb", "Mar", "Apr", "May", "Jun",
                          "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"),
               minor_breaks = NULL) + 
  scale_y_continuous(labels = numform::ff_num(zero = 0, digits = 2),
                     limits = c(-.3, 0)) +
  theme(legend.position = "bottom",
        plot.title = element_text(size = 20, hjust = 0.5),
        axis.title = element_text(size = 14),
        axis.text = element_text(size = 12),
        panel.grid.major = element_blank(), 
        panel.grid.minor = element_blank(),
        legend.text = element_text(size = 14),
        strip.text = element_text(size = 20),
        strip.background = element_blank())

p_con <- ggplot(subset(p_dat_wIdeo, ideo == "Conservative"), aes(x = week, y = pred, group = year)) +
  geom_vline(color = "red", aes(xintercept = as.Date("2020-5-22"))) +
  geom_linerange(aes(ymin = pred - 1.39*se,
                     ymax = pred + 1.39*se,
                     color = year),
                 position = position_dodge(width = .5),
                 lwd = .5) +
  geom_point(size = 2, aes(shape = year, color = year),
             position = position_dodge(width = .5)) +
  geom_smooth(data = subset(p_dat_wIdeo, ideo == "Conservative" & week < "2020-5-25"), 
              span = .8,
              aes(x = week, y = pred, color = year, linetype = year), se = FALSE) +
  geom_smooth(data = subset(p_dat_wIdeo, ideo == "Conservative" & week >= "2020-5-25"), 
              span = .8,
              aes(x = week, y = pred, color = year, linetype = year), se = FALSE) +
  theme_bw() +
  scale_linetype_manual(values = c("dotdash", "solid"), guide = FALSE) +
  scale_color_manual(values = colors,name = "") +
  scale_shape_manual(values = c(1, 16), name = "") +
  labs(y = "Predicted D-Scores",
       x = "", title = "Among White Conservatives") +
  scale_x_date(breaks = unique(subset(p_dat_wIdeo, ideo == "Conservative")$week)[c(1, 6, 10, 15, 19, 23, 28, 32, 37, 41, 45, 50)],
               labels = c("Jan", "Feb", "Mar", "Apr", "May", "Jun",
                          "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"),
               minor_breaks = NULL) + 
  scale_y_continuous(labels = numform::ff_num(zero = 0, digits = 2),
                     limits = c(.2, .59)) +
  theme(legend.position = "bottom",
        plot.title = element_text(size = 20, hjust = 0.5),
        axis.title = element_text(size = 14),
        axis.text = element_text(size = 12),
        panel.grid.major = element_blank(), 
        panel.grid.minor = element_blank(),
        legend.text = element_text(size = 14),
        strip.text = element_text(size = 20),
        strip.background = element_blank())

p_mod <- ggplot(subset(p_dat_wIdeo, ideo == "Moderate"), aes(x = week, y = pred, group = year)) +
  geom_vline(color = "red", aes(xintercept = as.Date("2020-5-22"))) +
  geom_linerange(aes(ymin = pred - 1.39*se,
                     ymax = pred + 1.39*se,
                     color = year),
                 position = position_dodge(width = .5),
                 lwd = .5) +
  geom_point(size = 2, aes(shape = year, color = year),
             position = position_dodge(width = .5)) +
  geom_smooth(data = subset(p_dat_wIdeo, ideo == "Moderate" & week < "2020-5-25"), 
              span = .8,
              aes(x = week, y = pred, color = year, linetype = year), se = FALSE) +
  geom_smooth(data = subset(p_dat_wIdeo, ideo == "Moderate" & week >= "2020-5-25"), 
              span = .8,
              aes(x = week, y = pred, color = year, linetype = year), se = FALSE) +
  theme_bw() +
  scale_linetype_manual(values = c("dotdash", "solid"), guide = FALSE) +
  scale_color_manual(values = colors,name = "") +
  scale_shape_manual(values = c(1, 16), name = "") +
  labs(y = "Predicted D-Scores",
       x = "", title = "Among White Neutrals\n(Moderates)") +
  scale_x_date(breaks = unique(subset(p_dat_wIdeo, ideo == "Moderate")$week)[c(1, 6, 10, 15, 19, 23, 28, 32, 37, 41, 45, 50)],
               labels = c("Jan", "Feb", "Mar", "Apr", "May", "Jun",
                          "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"),
               minor_breaks = NULL) + 
  scale_y_continuous(labels = numform::ff_num(zero = 0, digits = 2),
                     limits = c(.2, .5)) +
  theme(legend.position = "bottom",
        plot.title = element_text(size = 20, hjust = 0.5),
        axis.title = element_text(size = 14),
        axis.text = element_text(size = 12),
        panel.grid.major = element_blank(), 
        panel.grid.minor = element_blank(),
        legend.text = element_text(size = 14),
        strip.text = element_text(size = 20),
        strip.background = element_blank())

p_lib <- ggplot(subset(p_dat_wIdeo, ideo == "Liberal"), aes(x = week, y = pred, group = year)) +
  geom_vline(color = "red", aes(xintercept = as.Date("2020-5-22"))) +
  geom_linerange(aes(ymin = pred - 1.39*se,
                     ymax = pred + 1.39*se,
                     color = year),
                 position = position_dodge(width = .5),
                 lwd = .5) +
  geom_point(size = 2, aes(shape = year, color = year),
             position = position_dodge(width = .5)) +
  geom_smooth(data = subset(p_dat_wIdeo, ideo == "Liberal" & week < "2020-5-25"), 
              span = .8,
              aes(x = week, y = pred, color = year, linetype = year), se = FALSE) +
  geom_smooth(data = subset(p_dat_wIdeo, ideo == "Liberal" & week >= "2020-5-25"), 
              span = .8,
              aes(x = week, y = pred, color = year, linetype = year), se = FALSE) +
  theme_bw() +
  scale_linetype_manual(values = c("dotdash", "solid"), guide = FALSE) +
  scale_color_manual(values = colors,name = "") +
  scale_shape_manual(values = c(1, 16), name = "") +
  labs(y = "Predicted D-Scores",
       x = "", title = "Among White Liberals") +
  scale_x_date(breaks = unique(subset(p_dat_wIdeo, ideo == "Liberal")$week)[c(1, 6, 10, 15, 19, 23, 28, 32, 37, 41, 45, 50)],
               labels = c("Jan", "Feb", "Mar", "Apr", "May", "Jun",
                          "Jul", "Aug", "Sep", "Oct", "Nov", "Dec"),
               minor_breaks = NULL) + 
  scale_y_continuous(labels = numform::ff_num(zero = 0, digits = 2),
                     limits = c(.2, .5)) +
  theme(legend.position = "bottom",
        plot.title = element_text(size = 20, hjust = 0.5),
        axis.title = element_text(size = 14),
        axis.text = element_text(size = 12),
        panel.grid.major = element_blank(), 
        panel.grid.minor = element_blank(),
        legend.text = element_text(size = 14),
        strip.text = element_text(size = 20),
        strip.background = element_blank())



p_full + p_wht + p_blk + p_con + p_mod + p_lib +
  plot_layout(guides = 'collect', nrow = 2) + 
  plot_annotation(caption = "Note: Y-axis range for Black respondents is different due to different attitude levels. Predicted weekly means with 84% confidence intervals.") &  
  theme(plot.tag = element_text(size = 18),
        plot.caption = element_text(size = 14, hjust = 0),
        strip.text = element_text(size = 16),
        axis.title = element_text(size = 16),
        axis.text.x = element_text(size = 14, angle = 45, vjust = .8),
        panel.grid.major = element_blank(), 
        panel.grid.minor = element_blank(),
        legend.text = element_text(size = 16),
        legend.position = "bottom")

# Exporting to Stata ------------------------------------------------------
full_df <- p_full$data
wht_df <- p_wht$data[, c("year", "week", "pred", "se")]
blk_df <- p_blk$data[, c("year", "week", "pred", "se")]
con_df <- p_con$data[, c("year", "week", "pred", "se")]
mod_df <- p_mod$data[, c("year", "week", "pred", "se")]
lib_df <- p_lib$data[, c("year", "week", "pred", "se")]
full_df$df <- "full"
wht_df$df <- "white"
blk_df$df <- "black"
con_df$df <- "conservative"
mod_df$df <- "moderate"
lib_df$df <- "liberal"
dat_iat <- rbind(full_df, wht_df, blk_df, con_df, mod_df, lib_df)

dat_iat2 <- reshape(dat_iat, idvar = c("df", "week"), timevar = "year", direction = "wide")

write.dta(dat_iat2, file = "./data/iat_predictions.dta")

# End log file for script -------------------------------------------------
TeachingDemos::txtStop()